﻿using Microsoft.AspNetCore.Mvc;

namespace OpenAIProxy;

public class EventStreamResponse : IActionResult
{
    private readonly HttpResponseMessage _response;
    private readonly ITokenAccumulator _tokenAccumulator;
    private readonly ChatRecord _record;
    public decimal CostPerToken { get; set; }

    public EventStreamResponse(HttpResponseMessage response, ITokenAccumulator tokenAccumulator, ChatRecord record,decimal costPerToken)
    {
        _response = response;
        _tokenAccumulator = tokenAccumulator;
        _record = record;
        CostPerToken = costPerToken;
    }

    public async Task ExecuteResultAsync(ActionContext context)
    {
        var httpContext = context.HttpContext;
        var response = httpContext.Response;

        response.Headers.Add("Content-Type", "text/event-stream");
        response.Headers.Add("Cache-Control", "no-cache");
        response.Headers.Add("Connection", "keep-alive");

        using var stream = await _response.Content.ReadAsStreamAsync();
        using var reader = new StreamReader(stream);
        
        while (!reader.EndOfStream)
        {
            var line = await reader.ReadLineAsync();
            if (string.IsNullOrWhiteSpace(line)) continue;
            
            if (line.StartsWith("data: {"))
            {
                await response.WriteAsync($"{line}\n\n");
                await response.Body.FlushAsync();
                
                _tokenAccumulator.Append(line);
            }
            else if (line == "data: [DONE]")
            {
                await response.WriteAsync("data: [DONE]\n\n");
                await response.Body.FlushAsync();
                break;
            }
        }

        var (content, tokenCount) = _tokenAccumulator.GetAccumulatedResponse();
        _record.Response = content;
        _record.CompletionTokens = tokenCount;
    }
}
